{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#闪电模拟"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#1.0 闪电模拟+电击大楼\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import copy \n",
    "import seaborn as sns\n",
    "class lightning(object):\n",
    "    def __init__(self,f_size=[100,100],k=1):\n",
    "        self.f=np.zeros([100,100])\n",
    "        self.l=np.zeros([100,100])\n",
    "        self.l_group=[]\n",
    "        self.q_group=[]\n",
    "        self.g_group=[]\n",
    "        self.k=k\n",
    "        self.small_d=0.1\n",
    "        self.phi_group=[]\n",
    "    def ini_group_g_c(self):#闪电模拟\n",
    "        self.q_group=[]\n",
    "        for i in range(100):\n",
    "                self.q_group.append([int(len(self.f))-1,i])\n",
    "        self.l_group=[[int(len(self.f)-1),int(len(self.f[0])/2)]]#self.q_group#\n",
    "        self.g_group=[]\n",
    "        for i in range(100):\n",
    "            self.g_group.append([0,i])\n",
    "        #for i in range(43,58):\n",
    "            #self.g_group.append([30,i])\n",
    "        #for i in range(1,40):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.4)])\n",
    "        #for i in range(1,40):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.6)])\n",
    "        #for i in range(1,30):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.43)])\n",
    "        #for i in range(1,30):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.58)])\n",
    "        return 0\n",
    "    def ini_group_g_t(self):#电击大楼\n",
    "        self.q_group=[]\n",
    "        for i in range(100):\n",
    "                self.q_group.append([int(len(self.f))-1,i])\n",
    "        self.l_group=[[int(len(self.f)-1),int(len(self.f[0])/2)]]\n",
    "        self.g_group=[]\n",
    "        for i in range(100):\n",
    "            self.g_group.append([0,i])\n",
    "        for i in range(43,58):\n",
    "            self.g_group.append([30,i])\n",
    "        for i in range(1,40):\n",
    "            self.g_group.append([i,int(len(self.f)*0.4)])\n",
    "        for i in range(1,40):\n",
    "            self.g_group.append([i,int(len(self.f)*0.6)])\n",
    "        for i in range(1,30):\n",
    "            self.g_group.append([i,int(len(self.f)*0.43)])\n",
    "        for i in range(1,30):\n",
    "            self.g_group.append([i,int(len(self.f)*0.58)])\n",
    "        return 0\n",
    "    def test(self,time=1000,plt_show=True):\n",
    "        self.ini_group_g_c()#此处更改初始条件，选择模拟闪电or电击大楼\n",
    "        for t in range(time):          \n",
    "            self.l_refresh()\n",
    "            self.f_refresh()\n",
    "            self.next_step()\n",
    "        #for i in range(len(self.f)):\n",
    "            #self.g_group.append([len(self.f),1])\n",
    "            #print(self.l_group)\n",
    "            if plt_show and t%1==0:\n",
    "                fig=plt.figure(figsize=(10,10))\n",
    "                ax = sns.heatmap(self.f,cmap=\"binary_r\",cbar=False)\n",
    "                plt.xlim(0,len(self.f[0])-1)\n",
    "                plt.ylim(0,len(self.f[1])-1)\n",
    "                #for i in self.l_group:\n",
    "                    #plt.scatter(i[1],i[0],alpha=0.5,color=\"blue\")\n",
    "                plt.savefig(\"D:/3_bodies/DLA_pic/light_\"+str(t)+\".jpg\")\n",
    "                plt.show()\n",
    "        return \"done\"\n",
    "    def near(self):\n",
    "        near_group=[]\n",
    "        for i in range(len(self.f)):\n",
    "            for j in range(len(self.f[0])):\n",
    "                if i==0 and j!=0 and j!=len(self.l[0])-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif i==len(self.l) and j!=0 and j!=len(self.l[0])-1 and self.l[i,j]==0 and (self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i-1,j-1]==1 or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j==len(self.l[0]) and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j-1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j==0 and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i,j+1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j!=0 and j!=len(self.l[0])-1 and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i-1,j-1]==1 or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "        return near_group\n",
    "    def near_end(self):\n",
    "        near_group=[]\n",
    "        for i in [-1,0,1]:\n",
    "            for j in [-1,0,1]:\n",
    "                if (i!=0 or j!=0) and (self.l_group[-1][0]+i<0 and self.l_group[-1][0]+i>len(self.f) and\n",
    "                                      self.l_group[-1][1]+j<0 and self.l_group[-1][1]+j>len(self.f)):\n",
    "                    near_group.append([self.l_group[-1][0]+i,self.l_group[-1][1]+j])\n",
    "        return near_group\n",
    "    def l_refresh(self):\n",
    "        for i in self.l_group:\n",
    "            self.l[i[0],i[1]]=1\n",
    "        return \"l refreshed\"\n",
    "    def f_refresh(self):\n",
    "        f_now=copy.deepcopy(self.f)\n",
    "        for i in range(len(self.f)):\n",
    "            for j in range(len(self.f[0])):\n",
    "                f=0\n",
    "                for m in self.q_group:\n",
    "                    if i-m[0]!=0 and j-m[1]!=0:\n",
    "                    #-kq/r2\n",
    "                        f+=self.k/((i-m[0])**2+(j-m[1])**2)**0.5\n",
    "                    elif i-m[0]==0 and j-m[1]!=0:  \n",
    "                        f+=self.k/((j-m[1])**2)**0.5\n",
    "                    elif j-m[1]==0 and i-m[0]!=0:\n",
    "                        f+=self.k/((i-m[0])**2)**0.5\n",
    "                    else: \n",
    "                        f+=self.k/(self.small_d)\n",
    "                for m in self.g_group:\n",
    "                    if i-m[0]!=0 and j-m[1]!=0:\n",
    "                    #-kq/r2\n",
    "                        f-=self.k*2/((i-m[0])**2+(j-m[1])**2)**0.5\n",
    "                    elif i-m[0]==0 and j-m[1]!=0:  \n",
    "                        f-=self.k*2/((j-m[1])**2)**0.5\n",
    "                    elif j-m[1]==0 and i-m[0]!=0:\n",
    "                        f-=self.k*2/((i-m[0])**2)**0.5\n",
    "                    else: \n",
    "                        f-=self.k*2/(self.small_d)                \n",
    "                self.f[i,j]=f\n",
    "        return None\n",
    "    def grad(self,point):\n",
    "        grad_p=0\n",
    "        grad_p=np.max(self.f)-self.f[point[0],point[1]]\n",
    "        #for i in [-1,0,1]:\n",
    "            #for j in [-1,0,1]:\n",
    "                #if point[0]+i>0 and point[0]+i<len(self.f) and point[1]+i>0 and point[1]+i<len(self.f):\n",
    "                    #if abs()>grad_p:#abs(self.f[point[0],point[1]]-self.f[point[0]+i,point[1]+j])>grad_p:\n",
    "                        #grad_p=abs(self.f[point[0],point[1]]-self.f[point[0]+i,point[1]+j])\n",
    "        return grad_p**18\n",
    "    def get_phi(self,near_group):\n",
    "        self.phi_group=np.zeros(len(near_group))\n",
    "        for i in range(len(near_group)):\n",
    "            for j in range(len(self.l_group)):\n",
    "                dist=((self.l_group[j,0]-self.near_group[i,0])**2+(self.l_group[j,1]-self.near_group[i,1])**2)**0.5\n",
    "                self.phi_group[i]+=1/(4*math.pi*1*dist)\n",
    "        return \"phi_refresh\"\n",
    "    def next_step(self):\n",
    "        random_point=np.random.rand()\n",
    "        near_group=self.near()\n",
    "        #print(near_group)\n",
    "        grad_line=[]\n",
    "        for m in near_group:\n",
    "            #print(self.grad(m))\n",
    "            grad_line.append(self.grad(m))\n",
    "        grad_line/=np.sum(grad_line)\n",
    "        #print(random_point)\n",
    "        #print(grad_line)\n",
    "        for i in range(len(grad_line)):\n",
    "            random_point-=grad_line[i]\n",
    "            #print(random_point)\n",
    "            if random_point<0:\n",
    "                #print(i,grad_line[i])\n",
    "                self.l[near_group[i][0],near_group[i][1]]=1\n",
    "                self.l_group.append([near_group[i][0],near_group[i][1]])\n",
    "                self.q_group.append([near_group[i][0],near_group[i][1]])\n",
    "                return 0\n",
    "light_1=lightning()\n",
    "light_1.test()     \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#imgs2video\n",
    "import os\n",
    "import cv2\n",
    "from PIL import Image\n",
    "fourcc = cv2.VideoWriter_fourcc('X', 'V', 'I', 'D') \n",
    "path=r\"D:/3_bodies/DLA_pic/light_\"#input(\"path?\")\n",
    "save_path=r\"D:/3_bodies/light_tesra_3.mp4\"#input(\"save_path?\")\n",
    "post_video=cv2.VideoWriter(save_path, fourcc, 30,(720,720))\n",
    "for i in range(100):\n",
    "    print(path+str(i*10)+\".jpg\")\n",
    "    img=cv2.imread(path+str(i*10)+\".jpg\")\n",
    "    post_video.write(img)\n",
    "    del img\n",
    "    #if filename==\"0.jpg\":\n",
    "        #Image.fromarray(img).show()\n",
    "\n",
    "post_video.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Jacob ladder\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import copy \n",
    "import seaborn as sns\n",
    "class lightning_g(object):\n",
    "    def __init__(self,f_size=[100,100],k=1):\n",
    "        self.f=np.zeros([100,100])\n",
    "        self.l=np.zeros([100,100])\n",
    "        self.l_group=[]\n",
    "        self.q_group=[]\n",
    "        self.g_group=[]\n",
    "        self.k=k\n",
    "        self.small_d=0.1\n",
    "        self.phi_group=[]\n",
    "    def test(self,time=1000,plt_show=True):\n",
    "        self.q_group=[]\n",
    "        for i in range(100):\n",
    "            self.q_group.append([99-i,int(len(self.f)-0.4*i)-1])\n",
    "        self.l_group=self.q_group\n",
    "        self.g_group=[]\n",
    "        for i in range(100):\n",
    "            self.g_group.append([99-i,int(0.4*i)])\n",
    "        for t in range(time):\n",
    "            for q in self.q_group:\n",
    "                if q in self.g_group:\n",
    "                    del q\n",
    "            self.l_refresh()\n",
    "            self.f_refresh()\n",
    "            self.next_step()\n",
    "        #for i in range(len(self.f)):\n",
    "            #self.g_group.append([len(self.f),1])\n",
    "            #print(self.l_group)\n",
    "            if plt_show and t%1==0:\n",
    "                fig=plt.figure(figsize=(10,10))\n",
    "                ax = sns.heatmap(self.f,cmap=\"binary_r\",cbar=False)\n",
    "                plt.xlim(0,len(self.f[0])-1)\n",
    "                plt.ylim(0,len(self.f[1])-1)\n",
    "                #for i in self.l_group:\n",
    "                    #plt.scatter(i[1],i[0],alpha=0.5,color=\"blue\")\n",
    "                plt.savefig(\"D:/3_bodies/DLA_pic/single_\"+str(t)+\".jpg\")\n",
    "                plt.show()\n",
    "        return \"done\"\n",
    "    def near(self):\n",
    "        near_group=[]\n",
    "        for i in range(len(self.f)):\n",
    "            for j in range(len(self.f[0])):\n",
    "                if i==0 and j!=0 and j!=len(self.l[0])-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif i==len(self.l) and j!=0 and j!=len(self.l[0])-1 and self.l[i,j]==0 and (self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i-1,j-1]==1 or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j==len(self.l[0]) and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j-1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j==0 and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i,j+1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j!=0 and j!=len(self.l[0])-1 and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i-1,j-1]==1 or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "        return near_group\n",
    "    def near_end(self):\n",
    "        near_group=[]\n",
    "        for i in [-1,0,1]:\n",
    "            for j in [-1,0,1]:\n",
    "                if (i!=0 or j!=0) and (self.l_group[-1][0]+i<0 and self.l_group[-1][0]+i>len(self.f) and\n",
    "                                      self.l_group[-1][1]+j<0 and self.l_group[-1][1]+j>len(self.f)):\n",
    "                    near_group.append([self.l_group[-1][0]+i,self.l_group[-1][1]+j])\n",
    "        return near_group\n",
    "    def l_refresh(self):\n",
    "        for i in self.l_group:\n",
    "            self.l[i[0],i[1]]=1\n",
    "        return \"l refreshed\"\n",
    "    def f_refresh(self):\n",
    "        f_now=copy.deepcopy(self.f)\n",
    "        for i in range(len(self.f)):\n",
    "            for j in range(len(self.f[0])):\n",
    "                f=0\n",
    "                for m in self.q_group:\n",
    "                    if i-m[0]!=0 and j-m[1]!=0:\n",
    "                    #-kq/r2\n",
    "                        f+=self.k/((i-m[0])**2+(j-m[1])**2)**0.5\n",
    "                    elif i-m[0]==0 and j-m[1]!=0:  \n",
    "                        f+=self.k/((j-m[1])**2)**0.5\n",
    "                    elif j-m[1]==0 and i-m[0]!=0:\n",
    "                        f+=self.k/((i-m[0])**2)**0.5\n",
    "                    else: \n",
    "                        f+=self.k/(self.small_d)\n",
    "                for m in self.g_group:\n",
    "                    if i-m[0]!=0 and j-m[1]!=0:\n",
    "                    #-kq/r2\n",
    "                        f-=self.k*2/((i-m[0])**2+(j-m[1])**2)**0.5\n",
    "                    elif i-m[0]==0 and j-m[1]!=0:  \n",
    "                        f-=self.k*2/((j-m[1])**2)**0.5\n",
    "                    elif j-m[1]==0 and i-m[0]!=0:\n",
    "                        f-=self.k*2/((i-m[0])**2)**0.5\n",
    "                    else: \n",
    "                        f-=self.k*2/(self.small_d)                \n",
    "                self.f[i,j]=f\n",
    "        return None\n",
    "    def grad(self,point):\n",
    "        grad_p=0\n",
    "        grad_p=np.max(self.f)-self.f[point[0],point[1]]\n",
    "        #for i in [-1,0,1]:\n",
    "            #for j in [-1,0,1]:\n",
    "                #if point[0]+i>0 and point[0]+i<len(self.f) and point[1]+i>0 and point[1]+i<len(self.f):\n",
    "                    #if abs()>grad_p:#abs(self.f[point[0],point[1]]-self.f[point[0]+i,point[1]+j])>grad_p:\n",
    "                        #grad_p=abs(self.f[point[0],point[1]]-self.f[point[0]+i,point[1]+j])\n",
    "        return grad_p**18\n",
    "    def get_phi(self,near_group):\n",
    "        self.phi_group=np.zeros(len(near_group))\n",
    "        for i in range(len(near_group)):\n",
    "            for j in range(len(self.l_group)):\n",
    "                dist=((self.l_group[j,0]-self.near_group[i,0])**2+(self.l_group[j,1]-self.near_group[i,1])**2)**0.5\n",
    "                self.phi_group[i]+=1/(4*math.pi*1*dist)\n",
    "        return \"phi_refresh\"\n",
    "    def q_up(self):\n",
    "        for q in self.q_group\n",
    "    def next_step(self):\n",
    "        random_point=np.random.rand()\n",
    "        near_group=self.near()\n",
    "        #print(near_group)\n",
    "        grad_line=[]\n",
    "        for m in near_group:\n",
    "            #print(self.grad(m))\n",
    "            grad_line.append(self.grad(m))\n",
    "        grad_line/=np.sum(grad_line)\n",
    "        #print(random_point)\n",
    "        #print(grad_line)\n",
    "        for i in range(len(grad_line)):\n",
    "            random_point-=grad_line[i]\n",
    "            #print(random_point)\n",
    "            if random_point<0:\n",
    "                #print(i,grad_line[i])\n",
    "                self.l[near_group[i][0],near_group[i][1]]=1\n",
    "                self.l_group.append([near_group[i][0],near_group[i][1]])\n",
    "                self.q_group.append([near_group[i][0],near_group[i][1]])\n",
    "                return 0\n",
    "light_g_1=lightning_g()\n",
    "light_g_1.test()     \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#雷击任意物体，img转2值图，置于中央\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import copy \n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "class lightning(object):\n",
    "    def __init__(self,g_arr,f_size=[100,100],k=1):\n",
    "        self.f=np.zeros(f_size)\n",
    "        self.l=np.zeros(f_size)\n",
    "        self.l_group=[]\n",
    "        self.q_group=[]\n",
    "        self.g_group=[]\n",
    "        self.k=k\n",
    "        self.small_d=0.1\n",
    "        self.phi_group=[]\n",
    "        self.g_arr=g_arr\n",
    "    def ini_group_g_c(self):\n",
    "        self.q_group=[]\n",
    "        for i in range(100):\n",
    "                self.q_group.append([int(len(self.f))-1,i])\n",
    "        self.l_group=self.q_group#[[int(len(self.f)-1),int(len(self.f[0])/2)]]\n",
    "        self.g_group=[]\n",
    "        for i in range(100):\n",
    "            self.g_group.append([0,i])\n",
    "        #for i in range(43,58):\n",
    "            #self.g_group.append([30,i])\n",
    "        #for i in range(1,40):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.4)])\n",
    "        #for i in range(1,40):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.6)])\n",
    "        #for i in range(1,30):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.43)])\n",
    "        #for i in range(1,30):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.58)])\n",
    "        return 0\n",
    "    def get_g_from_img(self):\n",
    "        self.q_group=[]\n",
    "        for i in range(len(self.f)):\n",
    "                self.q_group.append([i,0])\n",
    "        self.l_group=[[int(len(self.f[0])/2),0]]\n",
    "        self.g_group=[]\n",
    "        for i in range(len(self.g_arr)):\n",
    "            for j in range(len(self.g_arr[0])):\n",
    "                if self.g_arr[i,j].all()==0:\n",
    "                    self.g_group.append([(len(self.f)+len(g_arr))/2-i,(len(self.f[0])-len(g_arr[0]))/2+j])\n",
    "        return 0\n",
    "    def ini_group_g_t(self):\n",
    "        self.q_group=[]\n",
    "        for i in range(100):\n",
    "                self.q_group.append([int(len(self.f))-1,i])\n",
    "        self.l_group=[[int(len(self.f)-1),int(len(self.f[0])/2)]]\n",
    "        self.g_group=[]\n",
    "        for i in range(100):\n",
    "            self.g_group.append([0,i])\n",
    "        for i in range(43,58):\n",
    "            self.g_group.append([30,i])\n",
    "        for i in range(1,40):\n",
    "            self.g_group.append([i,int(len(self.f)*0.4)])\n",
    "        for i in range(1,40):\n",
    "            self.g_group.append([i,int(len(self.f)*0.6)])\n",
    "        for i in range(1,30):\n",
    "            self.g_group.append([i,int(len(self.f)*0.43)])\n",
    "        for i in range(1,30):\n",
    "            self.g_group.append([i,int(len(self.f)*0.58)])\n",
    "        return 0\n",
    "    def test(self,time=1000,plt_show=True):\n",
    "        self.get_g_from_img()\n",
    "        for t in range(time):          \n",
    "            self.l_refresh()\n",
    "            self.f_refresh()\n",
    "            self.next_step()\n",
    "        #for i in range(len(self.f)):\n",
    "            #self.g_group.append([len(self.f),1])\n",
    "            #print(self.l_group)\n",
    "            if plt_show and t%1==0:\n",
    "                fig=plt.figure(figsize=(10,10))\n",
    "                ax = sns.heatmap(self.f,cmap=\"binary_r\",cbar=False)\n",
    "                plt.xlim(0,len(self.f[0])-1)\n",
    "                plt.ylim(0,len(self.f[1])-1)\n",
    "                #for i in self.l_group:\n",
    "                    #plt.scatter(i[1],i[0],alpha=0.5,color=\"blue\")\n",
    "                plt.savefig(\"D:/3_bodies/DLA_pic/light_\"+str(t)+\".jpg\")\n",
    "                plt.show()\n",
    "        return \"done\"\n",
    "    def near(self):\n",
    "        near_group=[]\n",
    "        for i in range(len(self.f)):\n",
    "            for j in range(len(self.f[0])):\n",
    "                if i==0 and j!=0 and j!=len(self.l[0])-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif i==len(self.l) and j!=0 and j!=len(self.l[0])-1 and self.l[i,j]==0 and (self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i-1,j-1]==1 or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j==len(self.l[0]) and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j-1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j==0 and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i,j+1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j!=0 and j!=len(self.l[0])-1 and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i-1,j-1]==1 or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "        return near_group\n",
    "    def near_end(self):\n",
    "        near_group=[]\n",
    "        for i in [-1,0,1]:\n",
    "            for j in [-1,0,1]:\n",
    "                if (i!=0 or j!=0) and (self.l_group[-1][0]+i<0 and self.l_group[-1][0]+i>len(self.f) and\n",
    "                                      self.l_group[-1][1]+j<0 and self.l_group[-1][1]+j>len(self.f)):\n",
    "                    near_group.append([self.l_group[-1][0]+i,self.l_group[-1][1]+j])\n",
    "        return near_group\n",
    "    def l_refresh(self):\n",
    "        for i in self.l_group:\n",
    "            self.l[i[0],i[1]]=1\n",
    "        return \"l refreshed\"\n",
    "    def f_refresh(self):\n",
    "        f_now=copy.deepcopy(self.f)\n",
    "        for i in range(len(self.f)):\n",
    "            for j in range(len(self.f[0])):\n",
    "                f=0\n",
    "                for m in self.q_group:\n",
    "                    if i-m[0]!=0 and j-m[1]!=0:\n",
    "                    #-kq/r2\n",
    "                        f+=self.k/((i-m[0])**2+(j-m[1])**2)**0.5\n",
    "                    elif i-m[0]==0 and j-m[1]!=0:  \n",
    "                        f+=self.k/((j-m[1])**2)**0.5\n",
    "                    elif j-m[1]==0 and i-m[0]!=0:\n",
    "                        f+=self.k/((i-m[0])**2)**0.5\n",
    "                    else: \n",
    "                        f+=self.k/(self.small_d)\n",
    "                for m in self.g_group:\n",
    "                    if i-m[0]!=0 and j-m[1]!=0:\n",
    "                    #-kq/r2\n",
    "                        f-=self.k*2/((i-m[0])**2+(j-m[1])**2)**0.5\n",
    "                    elif i-m[0]==0 and j-m[1]!=0:  \n",
    "                        f-=self.k*2/((j-m[1])**2)**0.5\n",
    "                    elif j-m[1]==0 and i-m[0]!=0:\n",
    "                        f-=self.k*2/((i-m[0])**2)**0.5\n",
    "                    else: \n",
    "                        f-=self.k*2/(self.small_d)                \n",
    "                self.f[i,j]=f\n",
    "        return None\n",
    "    def grad(self,point):\n",
    "        grad_p=0\n",
    "        grad_p=np.max(self.f)-self.f[point[0],point[1]]\n",
    "        #for i in [-1,0,1]:\n",
    "            #for j in [-1,0,1]:\n",
    "                #if point[0]+i>0 and point[0]+i<len(self.f) and point[1]+i>0 and point[1]+i<len(self.f):\n",
    "                    #if abs()>grad_p:#abs(self.f[point[0],point[1]]-self.f[point[0]+i,point[1]+j])>grad_p:\n",
    "                        #grad_p=abs(self.f[point[0],point[1]]-self.f[point[0]+i,point[1]+j])\n",
    "        return grad_p**18\n",
    "    def get_phi(self,near_group):\n",
    "        self.phi_group=np.zeros(len(near_group))\n",
    "        for i in range(len(near_group)):\n",
    "            for j in range(len(self.l_group)):\n",
    "                dist=((self.l_group[j,0]-self.near_group[i,0])**2+(self.l_group[j,1]-self.near_group[i,1])**2)**0.5\n",
    "                self.phi_group[i]+=1/(4*math.pi*1*dist)\n",
    "        return \"phi_refresh\"\n",
    "    def next_step(self):\n",
    "        random_point=np.random.rand()\n",
    "        near_group=self.near()\n",
    "        #print(near_group)\n",
    "        grad_line=[]\n",
    "        for m in near_group:\n",
    "            #print(self.grad(m))\n",
    "            grad_line.append(self.grad(m))\n",
    "        grad_line/=np.sum(grad_line)\n",
    "        #print(random_point)\n",
    "        #print(grad_line)\n",
    "        for i in range(len(grad_line)):\n",
    "            random_point-=grad_line[i]\n",
    "            #print(random_point)\n",
    "            if random_point<0:\n",
    "                #print(i,grad_line[i])\n",
    "                self.l[near_group[i][0],near_group[i][1]]=1\n",
    "                self.l_group.append([near_group[i][0],near_group[i][1]])\n",
    "                self.q_group.append([near_group[i][0],near_group[i][1]])\n",
    "                return 0\n",
    "def img2bi(img_path,size=[100,100],save_path=None):\n",
    "    img=Image.open(img_path)\n",
    "    img=img.resize(size)\n",
    "    arr=np.asarray(img)\n",
    "    ave=np.mean(arr)\n",
    "    print(arr.shape)\n",
    "    arr_bi=np.zeros(arr.shape)\n",
    "    for i in range(len(arr)):\n",
    "        for j in range(len(arr[i])):\n",
    "            if np.mean(arr[i,j])>ave:\n",
    "                arr_bi[i,j]=1\n",
    "    if save_path!=None:\n",
    "        arr_bi*=255\n",
    "        image=Image.fromarray(arr_bi.astype('uint8')).convert('RGB')\n",
    "        image.save(save_path)\n",
    "        image.show()\n",
    "    return arr_bi\n",
    "img_path=\"D:\\\\3_bodies\\\\DLA\\\\Poland.jpg\"\n",
    "g_arr=img2bi(img_path=img_path,size=[80,80])#\n",
    "print(g_arr.shape)\n",
    "light_1=lightning(g_arr=g_arr,f_size=[100,100])\n",
    "light_1.test(time=10000)     \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#特斯拉\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import copy \n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "class lightning(object):\n",
    "    def __init__(self,g_arr,f_size=[100,100],k=1):\n",
    "        self.f=np.zeros(f_size)\n",
    "        self.l=np.zeros(f_size)\n",
    "        self.l_group=[]\n",
    "        self.q_group=[]\n",
    "        self.g_group=[]\n",
    "        self.k=k\n",
    "        self.small_d=0.1\n",
    "        self.phi_group=[]\n",
    "        self.g_arr=g_arr\n",
    "    def ini_group_g_c(self):\n",
    "        self.q_group=[]\n",
    "        for i in range(100):\n",
    "                self.q_group.append([int(len(self.f))-1,i])\n",
    "        self.l_group=self.q_group#[[int(len(self.f)-1),int(len(self.f[0])/2)]]\n",
    "        self.g_group=[]\n",
    "        for i in range(100):\n",
    "            self.g_group.append([0,i])\n",
    "        #for i in range(43,58):\n",
    "            #self.g_group.append([30,i])\n",
    "        #for i in range(1,40):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.4)])\n",
    "        #for i in range(1,40):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.6)])\n",
    "        #for i in range(1,30):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.43)])\n",
    "        #for i in range(1,30):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.58)])\n",
    "        return 0\n",
    "    def get_g_from_img(self):\n",
    "        self.q_group=[]\n",
    "        for i in range(len(self.f)):\n",
    "                self.q_group.append([int(len(self.f))-1,i])\n",
    "        self.l_group=[[int(len(self.f)-1),int(len(self.f[0])/2)]]\n",
    "        self.g_group=[]\n",
    "        for i in range(len(self.g_arr)):\n",
    "            for j in range(len(self.g_arr[0])):\n",
    "                if self.g_arr[i,j].all()==1:\n",
    "                    self.g_group.append([(len(self.f)+len(g_arr))/2-i,(len(self.f[0])-len(g_arr[0]))/2+j])\n",
    "        return 0\n",
    "    def ini_group_g_t(self):\n",
    "        self.q_group=[]\n",
    "        for i in range(100):\n",
    "                self.q_group.append([int(len(self.f))-1,i])\n",
    "        self.l_group=[[int(len(self.f)-1),int(len(self.f[0])/2)]]\n",
    "        self.g_group=[]\n",
    "        for i in range(100):\n",
    "            self.g_group.append([0,i])\n",
    "        for i in range(43,58):\n",
    "            self.g_group.append([30,i])\n",
    "        for i in range(1,40):\n",
    "            self.g_group.append([i,int(len(self.f)*0.4)])\n",
    "        for i in range(1,40):\n",
    "            self.g_group.append([i,int(len(self.f)*0.6)])\n",
    "        for i in range(1,30):\n",
    "            self.g_group.append([i,int(len(self.f)*0.43)])\n",
    "        for i in range(1,30):\n",
    "            self.g_group.append([i,int(len(self.f)*0.58)])\n",
    "        return 0\n",
    "    def grounded(self):\n",
    "        for g in self.g_group:\n",
    "            if g in self.q_group:\n",
    "                return True\n",
    "        return False\n",
    "    def ini_tesla_g(self):\n",
    "        self.q_group=[]\n",
    "        for i in range(1,int(len(self.f)/2)+1):\n",
    "                self.q_group.append([i,int(len(self.f)/2)])\n",
    "        self.l_group=[]\n",
    "        self.l=np.zeros(self.l.shape)\n",
    "        self.l_group=[[int(len(self.f[0])/2),int(len(self.f[0])/2)]]\n",
    "        self.g_group=[]\n",
    "        #for i in range(int(len(self.f)*2/3),len(self.f)):\n",
    "                #self.g_group.append([int(len(self.f)/2),i])\n",
    "        for i in range(len(self.f)):\n",
    "            self.g_group.append([0,i])\n",
    "    def ini_tesla_2g(self):\n",
    "        self.q_group=[]\n",
    "        for i in range(int(len(self.f)/3)+1):\n",
    "                self.q_group.append([i,int(len(self.f)/2)])\n",
    "        self.l_group=[]\n",
    "        self.l=np.zeros(self.l.shape)\n",
    "        self.l_group=[[int(len(self.f[0])/3),int(len(self.f[0])/2)]]\n",
    "        self.g_group=[]\n",
    "        for i in range(int(len(self.f)*2/3),len(self.f)):\n",
    "                self.g_group.append([i,int(len(self.f)/2)])\n",
    "        #for i in range(len(self.f)):\n",
    "            #self.g_group.append([0,i])\n",
    "        return 0\n",
    "    def test(self,time=1000,plt_show=True):\n",
    "        self.get_g_from_img()\n",
    "        for t in range(time):          \n",
    "            self.l_refresh()\n",
    "            self.f_refresh()\n",
    "            self.next_step()\n",
    "        #for i in range(len(self.f)):\n",
    "            #self.g_group.append([len(self.f),1])\n",
    "            #print(self.l_group)\n",
    "            if plt_show and t%100==0:\n",
    "                fig=plt.figure(figsize=(10,10))\n",
    "                ax = sns.heatmap(self.f,cmap=\"binary_r\",cbar=False)\n",
    "                plt.xlim(0,len(self.f[0])-1)\n",
    "                plt.ylim(0,len(self.f[1])-1)\n",
    "                #for i in self.l_group:\n",
    "                    #plt.scatter(i[1],i[0],alpha=0.5,color=\"blue\")\n",
    "                plt.savefig(\"D:/3_bodies/DLA_pic/light_\"+str(t)+\".jpg\")\n",
    "                plt.show()\n",
    "        return \"done\"\n",
    "    def test_spike(self,time=1000,period=100,plt_show=True):\n",
    "        for t in range(time): \n",
    "            self.ini_tesla_g()\n",
    "            for p in range(period):\n",
    "                self.l_refresh()\n",
    "                self.f_refresh()\n",
    "                self.next_step()\n",
    "        #for i in range(len(self.f)):\n",
    "            #self.g_group.append([len(self.f),1])\n",
    "            #print(self.l_group)\n",
    "                if plt_show and p%10==0:\n",
    "                    fig=plt.figure(figsize=(10,10))\n",
    "                    ax = sns.heatmap(self.f,cmap=\"binary_r\",cbar=False)\n",
    "                    plt.xlim(0,len(self.f[0])-1)\n",
    "                    plt.ylim(0,len(self.f[1])-1)\n",
    "                #for i in self.l_group:\n",
    "                    #plt.scatter(i[1],i[0],alpha=0.5,color=\"blue\")\n",
    "                    plt.savefig(\"D:/3_bodies/DLA_pic/light_\"+str(int(t*10+p/10))+\".jpg\")\n",
    "                    plt.show()\n",
    "    def test_spike_2(self,time=1000,plt_show=True):\n",
    "        self.ini_tesla_2g()\n",
    "        end_count=0\n",
    "        for t in range(time): \n",
    "            self.l_refresh()\n",
    "            self.f_refresh()\n",
    "            self.next_step()\n",
    "        #for i in range(len(self.f)):\n",
    "            #self.g_group.append([len(self.f),1])\n",
    "            #print(self.l_group)\n",
    "            if plt_show and t%10==0:\n",
    "                fig=plt.figure(figsize=(10,10))\n",
    "                ax = sns.heatmap(self.f,cmap=\"binary_r\",cbar=False)\n",
    "                plt.xlim(0,len(self.f[0])-1)\n",
    "                plt.ylim(0,len(self.f[1])-1)\n",
    "                #for i in self.l_group:\n",
    "                    #plt.scatter(i[1],i[0],alpha=0.5,color=\"blue\")\n",
    "                plt.savefig(\"D:/3_bodies/DLA_pic/light_\"+str(int(t/10))+\".jpg\")\n",
    "                plt.show()\n",
    "            if self.grounded():\n",
    "                end_count+=1\n",
    "            if end_count==10:\n",
    "                end_count=0\n",
    "                self.ini_tesla_2g()\n",
    "        return 0\n",
    "    def near(self):\n",
    "        near_group=[]\n",
    "        for i in range(len(self.f)):\n",
    "            for j in range(len(self.f[0])):\n",
    "                if i==0 and j!=0 and j!=len(self.l[0])-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif i==len(self.l) and j!=0 and j!=len(self.l[0])-1 and self.l[i,j]==0 and (self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i-1,j-1]==1 or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j==len(self.l[0]) and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j-1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j==0 and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i,j+1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j!=0 and j!=len(self.l[0])-1 and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i-1,j-1]==1 or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "        return near_group\n",
    "    def near_end(self):\n",
    "        near_group=[]\n",
    "        for i in [-1,0,1]:\n",
    "            for j in [-1,0,1]:\n",
    "                if (i!=0 or j!=0) and (self.l_group[-1][0]+i<0 and self.l_group[-1][0]+i>len(self.f) and\n",
    "                                      self.l_group[-1][1]+j<0 and self.l_group[-1][1]+j>len(self.f)):\n",
    "                    near_group.append([self.l_group[-1][0]+i,self.l_group[-1][1]+j])\n",
    "        return near_group\n",
    "    def l_refresh(self):\n",
    "        for i in self.l_group:\n",
    "            self.l[i[0],i[1]]=1\n",
    "        return \"l refreshed\"\n",
    "    def f_refresh(self):\n",
    "        f_now=copy.deepcopy(self.f)\n",
    "        for i in range(len(self.f)):\n",
    "            for j in range(len(self.f[0])):\n",
    "                f=0\n",
    "                for m in self.q_group:\n",
    "                    if i-m[0]!=0 or j-m[1]!=0:\n",
    "                    #-kq/r2\n",
    "                        f+=self.k/((i-m[0])**2+(j-m[1])**2)**0.5\n",
    "                    else: \n",
    "                        f+=self.k/(self.small_d)\n",
    "                for m in self.g_group:\n",
    "                    if i-m[0]!=0 or j-m[1]!=0:\n",
    "                    #-kq/r2\n",
    "                        f-=self.k*2/((i-m[0])**2+(j-m[1])**2)**0.5\n",
    "                    else: \n",
    "                        f-=self.k*2/(self.small_d)                \n",
    "                self.f[i,j]=f\n",
    "        return None\n",
    "    def grad(self,point):\n",
    "        grad_p=0\n",
    "        grad_p=np.max(self.f)-self.f[point[0],point[1]]\n",
    "        #for i in [-1,0,1]:\n",
    "            #for j in [-1,0,1]:\n",
    "                #if point[0]+i>0 and point[0]+i<len(self.f) and point[1]+i>0 and point[1]+i<len(self.f):\n",
    "                    #if abs()>grad_p:#abs(self.f[point[0],point[1]]-self.f[point[0]+i,point[1]+j])>grad_p:\n",
    "                        #grad_p=abs(self.f[point[0],point[1]]-self.f[point[0]+i,point[1]+j])\n",
    "        return grad_p**18\n",
    "    def get_phi(self,near_group):\n",
    "        self.phi_group=np.zeros(len(near_group))\n",
    "        for i in range(len(near_group)):\n",
    "            for j in range(len(self.l_group)):\n",
    "                dist=((self.l_group[j,0]-self.near_group[i,0])**2+(self.l_group[j,1]-self.near_group[i,1])**2)**0.5\n",
    "                self.phi_group[i]+=1/(4*math.pi*1*dist)\n",
    "        return \"phi_refresh\"\n",
    "    def next_step(self):\n",
    "        random_point=np.random.rand()\n",
    "        near_group=self.near()\n",
    "        #print(near_group)\n",
    "        grad_line=[]\n",
    "        for m in near_group:\n",
    "            #print(self.grad(m))\n",
    "            grad_line.append(self.grad(m))\n",
    "        grad_line/=np.sum(grad_line)\n",
    "        #print(random_point)\n",
    "        #print(grad_line)\n",
    "        for i in range(len(grad_line)):\n",
    "            random_point-=grad_line[i]\n",
    "            #print(random_point)\n",
    "            if random_point<0:\n",
    "                #print(i,grad_line[i])\n",
    "                #self.l[near_group[i][0],near_group[i][1]]=1\n",
    "                self.l_group.append([near_group[i][0],near_group[i][1]])\n",
    "                self.q_group.append([near_group[i][0],near_group[i][1]])\n",
    "                return 0\n",
    "def img2bi(img_path,size=[100,100],save_path=None):\n",
    "    img=Image.open(img_path)\n",
    "    img=img.resize(size)\n",
    "    arr=np.asarray(img)\n",
    "    ave=np.mean(arr)\n",
    "    print(arr.shape)\n",
    "    arr_bi=np.zeros(arr.shape)\n",
    "    for i in range(len(arr)):\n",
    "        for j in range(len(arr[i])):\n",
    "            if np.mean(arr[i,j])>ave:\n",
    "                arr_bi[i,j]=1\n",
    "    if save_path!=None:\n",
    "        arr_bi*=255\n",
    "        image=Image.fromarray(arr_bi.astype('uint8')).convert('RGB')\n",
    "        image.save(save_path)\n",
    "        image.show()\n",
    "    return arr_bi\n",
    "img_path=\"D:\\\\3_bodies\\\\DLA\\\\test2.jpg\"\n",
    "g_arr=img2bi(img_path=img_path,size=[80,50])#\n",
    "print(g_arr.shape)\n",
    "light_1=lightning(g_arr=g_arr,f_size=[100,100])\n",
    "light_1.test_spike_2(time=10000)     \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#贪婪，改进中\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import copy \n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "class lightning(object):\n",
    "    def __init__(self,g_arr,f_size=[100,100],k=1):\n",
    "        self.f=np.zeros(f_size)\n",
    "        self.l=np.zeros(f_size)\n",
    "        self.l_group=[]\n",
    "        self.q_group=[]\n",
    "        self.g_group=[]\n",
    "        self.k=k\n",
    "        self.small_d=0.1\n",
    "        self.phi_group=[]\n",
    "        self.g_arr=g_arr\n",
    "    def ini_group_g_c(self):\n",
    "        self.q_group=[]\n",
    "        for i in range(100):\n",
    "                self.q_group.append([int(len(self.f))-1,i])\n",
    "        self.l_group=self.q_group#[[int(len(self.f)-1),int(len(self.f[0])/2)]]\n",
    "        self.g_group=[]\n",
    "        for i in range(100):\n",
    "            self.g_group.append([0,i])\n",
    "        #for i in range(43,58):\n",
    "            #self.g_group.append([30,i])\n",
    "        #for i in range(1,40):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.4)])\n",
    "        #for i in range(1,40):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.6)])\n",
    "        #for i in range(1,30):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.43)])\n",
    "        #for i in range(1,30):\n",
    "            #self.g_group.append([i,int(len(self.f)*0.58)])\n",
    "        return 0\n",
    "    def get_g_from_img(self):\n",
    "        self.q_group=[]\n",
    "        for i in range(len(self.f)):\n",
    "                self.q_group.append([int(len(self.f))-1,i])\n",
    "        self.l_group=[[int(len(self.f)-1),int(len(self.f[0])/2)]]\n",
    "        self.g_group=[]\n",
    "        for i in range(len(self.g_arr)):\n",
    "            for j in range(len(self.g_arr[0])):\n",
    "                if self.g_arr[i,j].all()==1:\n",
    "                    self.g_group.append([(len(self.f)+len(g_arr))/2-i,(len(self.f[0])-len(g_arr[0]))/2+j])\n",
    "        return 0\n",
    "    def ini_group_g_t(self):\n",
    "        self.q_group=[]\n",
    "        for i in range(100):\n",
    "                self.q_group.append([int(len(self.f))-1,i])\n",
    "        self.l_group=[[int(len(self.f)-1),int(len(self.f[0])/2)]]\n",
    "        self.g_group=[]\n",
    "        for i in range(100):\n",
    "            self.g_group.append([0,i])\n",
    "        for i in range(43,58):\n",
    "            self.g_group.append([30,i])\n",
    "        for i in range(1,40):\n",
    "            self.g_group.append([i,int(len(self.f)*0.4)])\n",
    "        for i in range(1,40):\n",
    "            self.g_group.append([i,int(len(self.f)*0.6)])\n",
    "        for i in range(1,30):\n",
    "            self.g_group.append([i,int(len(self.f)*0.43)])\n",
    "        for i in range(1,30):\n",
    "            self.g_group.append([i,int(len(self.f)*0.58)])\n",
    "        return 0\n",
    "    def grounded(self):\n",
    "        for g in self.g_group:\n",
    "            if g in self.q_group:\n",
    "                return True\n",
    "        return False\n",
    "    def ini_tesla_g(self):\n",
    "        self.q_group=[]\n",
    "        for i in range(1,int(len(self.f)/2)+1):\n",
    "                self.q_group.append([i,int(len(self.f)/2)])\n",
    "        self.l_group=[]\n",
    "        self.l=np.zeros(self.l.shape)\n",
    "        self.l_group=[[int(len(self.f[0])/2),int(len(self.f[0])/2)]]\n",
    "        self.g_group=[]\n",
    "        #for i in range(int(len(self.f)*2/3),len(self.f)):\n",
    "                #self.g_group.append([int(len(self.f)/2),i])\n",
    "        for i in range(len(self.f)):\n",
    "            self.g_group.append([0,i])\n",
    "    def ini_tesla_2g(self):\n",
    "        self.q_group=[]\n",
    "        for i in range(int(len(self.f)/3)+1):\n",
    "                self.q_group.append([i,int(len(self.f)/2)])\n",
    "        self.l_group=[]\n",
    "        self.l=np.zeros(self.l.shape)\n",
    "        self.l_group=[[int(len(self.f[0])/3),int(len(self.f[0])/2)]]\n",
    "        self.g_group=[]\n",
    "        for i in range(int(len(self.f)*2/3),len(self.f)):\n",
    "                self.g_group.append([i,int(len(self.f)/2)])\n",
    "        #for i in range(len(self.f)):\n",
    "            #self.g_group.append([0,i])\n",
    "        return 0\n",
    "    def test(self,time=1000,plt_show=True):\n",
    "        self.get_g_from_img()\n",
    "        for t in range(time):          \n",
    "            self.l_refresh()\n",
    "            self.f_refresh()\n",
    "            self.next_step()\n",
    "        #for i in range(len(self.f)):\n",
    "            #self.g_group.append([len(self.f),1])\n",
    "            #print(self.l_group)\n",
    "            if plt_show and t%100==0:\n",
    "                fig=plt.figure(figsize=(10,10))\n",
    "                ax = sns.heatmap(self.f,cmap=\"binary_r\",cbar=False)\n",
    "                plt.xlim(0,len(self.f[0])-1)\n",
    "                plt.ylim(0,len(self.f[1])-1)\n",
    "                #for i in self.l_group:\n",
    "                    #plt.scatter(i[1],i[0],alpha=0.5,color=\"blue\")\n",
    "                plt.savefig(\"D:/3_bodies/DLA_pic/light_\"+str(t)+\".jpg\")\n",
    "                plt.show()\n",
    "        return \"done\"\n",
    "    def test_spike(self,time=1000,period=100,plt_show=True):\n",
    "        for t in range(time): \n",
    "            self.ini_tesla_g()\n",
    "            for p in range(period):\n",
    "                self.l_refresh()\n",
    "                self.f_refresh()\n",
    "                self.next_step()\n",
    "        #for i in range(len(self.f)):\n",
    "            #self.g_group.append([len(self.f),1])\n",
    "            #print(self.l_group)\n",
    "                if plt_show and p%10==0:\n",
    "                    fig=plt.figure(figsize=(10,10))\n",
    "                    ax = sns.heatmap(self.f,cmap=\"binary_r\",cbar=False)\n",
    "                    plt.xlim(0,len(self.f[0])-1)\n",
    "                    plt.ylim(0,len(self.f[1])-1)\n",
    "                #for i in self.l_group:\n",
    "                    #plt.scatter(i[1],i[0],alpha=0.5,color=\"blue\")\n",
    "                    plt.savefig(\"D:/3_bodies/DLA_pic/light_\"+str(int(t*10+p/10))+\".jpg\")\n",
    "                    plt.show()\n",
    "    def test_spike_2(self,time=1000,plt_show=True):\n",
    "        self.ini_tesla_2g()\n",
    "        end_count=0\n",
    "        for t in range(time): \n",
    "            self.l_refresh()\n",
    "            self.f_refresh()\n",
    "            self.next_step()\n",
    "        #for i in range(len(self.f)):\n",
    "            #self.g_group.append([len(self.f),1])\n",
    "            #print(self.l_group)\n",
    "            if plt_show and t%10==0:\n",
    "                fig=plt.figure(figsize=(10,10))\n",
    "                ax = sns.heatmap(self.f,cmap=\"binary_r\",cbar=False)\n",
    "                plt.xlim(0,len(self.f[0])-1)\n",
    "                plt.ylim(0,len(self.f[1])-1)\n",
    "                #for i in self.l_group:\n",
    "                    #plt.scatter(i[1],i[0],alpha=0.5,color=\"blue\")\n",
    "                plt.savefig(\"D:/3_bodies/DLA_pic/light_\"+str(int(t/10))+\".jpg\")\n",
    "                plt.show()\n",
    "            if self.grounded():\n",
    "                end_count+=1\n",
    "            if end_count==10:\n",
    "                end_count=0\n",
    "                self.ini_tesla_2g()\n",
    "        return 0\n",
    "    def near(self):\n",
    "        near_group=[]\n",
    "        for i in range(len(self.f)):\n",
    "            for j in range(len(self.f[0])):\n",
    "                if i==0 and j!=0 and j!=len(self.l[0])-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif i==len(self.l) and j!=0 and j!=len(self.l[0])-1 and self.l[i,j]==0 and (self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i-1,j-1]==1 or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j==len(self.l[0]) and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j-1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j==0 and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i,j+1]==1):\n",
    "                    near_group.append([i,j])\n",
    "                elif j!=0 and j!=len(self.l[0])-1 and i!=0 and i!=len(self.l)-1 and self.l[i,j]==0 and (self.l[i+1,j]==1 or self.l[i+1,j+1]==1 or self.l[i+1,j-1]==1 \n",
    "                or self.l[i-1,j]==1 or self.l[i-1,j+1]==1 or self.l[i-1,j-1]==1 or self.l[i,j+1]==1 or self.l[i,j-1]==1):\n",
    "                    near_group.append([i,j])\n",
    "        return near_group\n",
    "    def near_end(self):\n",
    "        near_group=[]\n",
    "        for i in [-1,0,1]:\n",
    "            for j in [-1,0,1]:\n",
    "                if (i!=0 or j!=0) and (self.l_group[-1][0]+i<0 and self.l_group[-1][0]+i>len(self.f) and\n",
    "                                      self.l_group[-1][1]+j<0 and self.l_group[-1][1]+j>len(self.f)):\n",
    "                    near_group.append([self.l_group[-1][0]+i,self.l_group[-1][1]+j])\n",
    "        return near_group\n",
    "    def l_refresh(self):\n",
    "        for i in self.l_group:\n",
    "            self.l[i[0],i[1]]=1\n",
    "        return \"l refreshed\"\n",
    "    def f_refresh(self):\n",
    "        f_now=copy.deepcopy(self.f)\n",
    "        for i in range(len(self.f)):\n",
    "            for j in range(len(self.f[0])):\n",
    "                f=0\n",
    "                for m in self.q_group:\n",
    "                    if i-m[0]!=0 or j-m[1]!=0:\n",
    "                    #-kq/r2\n",
    "                        f+=self.k/((i-m[0])**2+(j-m[1])**2)**0.5\n",
    "                    else: \n",
    "                        f+=self.k/(self.small_d)\n",
    "                for m in self.g_group:\n",
    "                    if i-m[0]!=0 or j-m[1]!=0:\n",
    "                    #-kq/r2\n",
    "                        f-=self.k*2/((i-m[0])**2+(j-m[1])**2)**0.5\n",
    "                    else: \n",
    "                        f-=self.k*2/(self.small_d)                \n",
    "                self.f[i,j]=f\n",
    "        return None\n",
    "    def grad(self,point):\n",
    "        grad_p=0\n",
    "        grad_p=np.max(self.f)-self.f[point[0],point[1]]\n",
    "        #for i in [-1,0,1]:\n",
    "            #for j in [-1,0,1]:\n",
    "                #if point[0]+i>0 and point[0]+i<len(self.f) and point[1]+i>0 and point[1]+i<len(self.f):\n",
    "                    #if abs()>grad_p:#abs(self.f[point[0],point[1]]-self.f[point[0]+i,point[1]+j])>grad_p:\n",
    "                        #grad_p=abs(self.f[point[0],point[1]]-self.f[point[0]+i,point[1]+j])\n",
    "        return grad_p\n",
    "    def get_phi(self,near_group):\n",
    "        self.phi_group=np.zeros(len(near_group))\n",
    "        for i in range(len(near_group)):\n",
    "            for j in range(len(self.l_group)):\n",
    "                dist=((self.l_group[j,0]-self.near_group[i,0])**2+(self.l_group[j,1]-self.near_group[i,1])**2)**0.5\n",
    "                self.phi_group[i]+=1/(4*math.pi*1*dist)\n",
    "        return \"phi_refresh\"\n",
    "    def next_step(self):\n",
    "        random_point=np.random.rand()\n",
    "        near_group=self.near()\n",
    "        #print(near_group)\n",
    "        grad_line=[]\n",
    "        for m in near_group:\n",
    "            #print(self.grad(m))\n",
    "            grad_line.append(self.grad(m))\n",
    "        grad_line=np.array(grad_line)\n",
    "        #print(random_point)\n",
    "        #print(grad_line)\n",
    "        self.l_group.append([near_group[np.argmax(grad_line)][0],near_group[np.argmax(grad_line)][1]])\n",
    "        self.q_group.append([near_group[np.argmax(grad_line)][0],near_group[np.argmax(grad_line)][1]])\n",
    "        return 0\n",
    "def img2bi(img_path,size=[100,100],save_path=None):\n",
    "    img=Image.open(img_path)\n",
    "    img=img.resize(size)\n",
    "    arr=np.asarray(img)\n",
    "    ave=np.mean(arr)\n",
    "    print(arr.shape)\n",
    "    arr_bi=np.zeros(arr.shape)\n",
    "    for i in range(len(arr)):\n",
    "        for j in range(len(arr[i])):\n",
    "            if np.mean(arr[i,j])>ave:\n",
    "                arr_bi[i,j]=1\n",
    "    if save_path!=None:\n",
    "        arr_bi*=255\n",
    "        image=Image.fromarray(arr_bi.astype('uint8')).convert('RGB')\n",
    "        image.save(save_path)\n",
    "        image.show()\n",
    "    return arr_bi\n",
    "img_path=\"D:\\\\3_bodies\\\\DLA\\\\test2.jpg\"\n",
    "g_arr=img2bi(img_path=img_path,size=[80,50])#\n",
    "print(g_arr.shape)\n",
    "light_1=lightning(g_arr=g_arr,f_size=[100,100])\n",
    "light_1.test_spike_2(time=10000)     \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5\n"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
